import timm
from dataclasses import dataclass
from torch.optim import *
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import Compose
from torch.nn import Module
@dataclass(repr=True)
class CLSConfig:
    image_size : int  # the generated image resolution
    train_batch_size : int
    eval_batch_size : int  # how many images to sample during evaluation
    num_epochs : int
    nc : int
    warm_epochs : int
    learning_rate : float
    momentum : float  # SGD momentum/Adam beta1
    weight_decay : float
    optimizer : SGD
    mixed_precision : str  # `no` for float32, `fp16` for automatic mixed precision
    output_dir : str  # the model namy locally and on the HF Hub
    seed : int
    model_args : dict
    model : Module
    index2label_file_path : str
    index2label : dict
    train_txt : str
    valid_txt : str
    test_txt : str
    dataset_path : str
    dataset_name : str
    train_images_file_list : list
    train_labels_list : list
    valid_images_file_list : list
    valid_labels_list : list
    test_images_file_list : list
    test_labels_list : list
    train_dataset : Dataset
    valid_dataset : Dataset
    train_dataloader : DataLoader
    valid_dataloader : DataLoader
    cache : bool
    workers : int
    mosaic : bool
    transforms : Compose
    max_nums : int
flower_clsconfig_dict = dict(
    image_size = 112,  # the generated image resolution
    train_batch_size = 32,
    eval_batch_size = 32,  # how many images to sample during evaluation
    num_epochs = 10,
    nc = 102,
    warm_epochs = 3,
    learning_rate = 1e-3,
    momentum = 0.937,  # SGD momentum/Adam beta1
    weight_decay =  0.00001,
    optimizer = AdamW,
    mixed_precision = 'fp16',  # `no` for float32, `fp16` for automatic mixed precision
    output_dir = 'cls' , # the model namy locally and on the HF Hub
    seed = 213,
    model_args = dict(model_type = 'mymodel', model_config = 'ylcls/configs/ResNet/ResNet50.yaml'),
    # model_args = dict(model_type='timmmodel', model_config=dict(model_name='resnet50', pretrained=False, num_classes=102)),
    model = None,
    index2label_file_path = 'oxford-102-flowers/index2label.json',
    index2label = None,
    train_txt = 'train1.txt',
    valid_txt = 'valid.txt',
    test_txt = 'test.txt',
    dataset_path = "oxford-102-flowers",
    dataset_name = 'CLSDataset',
    train_images_file_list = None,
    train_labels_list = None,
    valid_images_file_list = None,
    valid_labels_list = None,
    test_images_file_list = None,
    test_labels_list = None,
    train_dataset = None,
    valid_dataset = None,
    train_dataloader = None,
    valid_dataloader = None,
    cache = False,
    workers = 0,
    mosaic = True,
    transforms = None,
    max_nums = None
)
garbage_clsconfig_dict = dict(
    image_size = 224,  # the generated image resolution
    train_batch_size = 16,
    eval_batch_size = 16,  # how many images to sample during evaluation
    num_epochs = 10,
    nc = 40,
    warm_epochs = 3,
    learning_rate = 1e-3,
    momentum = 0.937,  # SGD momentum/Adam beta1
    weight_decay =  0.0001,
    optimizer = SGD,
    mixed_precision = 'fp16',  # `no` for float32, `fp16` for automatic mixed precision
    output_dir = 'cls' , # the model namy locally and on the HF Hub
    seed = 213,
    model_args = dict(model_type = 'mymodel', model_config = 'ylcls/configs/ResNet/ResNet50.yaml'),
    # model_args = dict(model_type='timmmodel', model_config=dict(model_name='resnet50', pretrained=False, num_classes=102)),
    model = None,
    index2label_file_path = 'garbage/index2label.json',
    index2label = None,
    train_txt='train.txt',
    valid_txt='valid.txt',
    test_txt='test.txt',
    dataset_path = "garbage",
    dataset_name = 'CLSDataset',
    train_images_file_list = None,
    train_labels_list = None,
    valid_images_file_list = None,
    valid_labels_list = None,
    test_images_file_list = None,
    test_labels_list = None,
    train_dataset = None,
    valid_dataset = None,
    train_dataloader = None,
    valid_dataloader = None,
    cache = False,
    workers = 0,
    mosaic = False,
    transforms = None,
    max_nums = None
)
clsconfig_dict = dict(
    image_size = 224,  # the generated image resolution
    train_batch_size = 16,
    eval_batch_size = 16,  # how many images to sample during evaluation
    num_epochs = 10,
    nc = 40,
    warm_epochs = 3,
    learning_rate = 1e-3,
    momentum = 0.937,  # SGD momentum/Adam beta1
    weight_decay =  0.0001,
    optimizer = SGD,
    mixed_precision = 'fp16',  # `no` for float32, `fp16` for automatic mixed precision
    output_dir = 'cls' , # the model namy locally and on the HF Hub
    seed = 213,
    model_args = dict(model_type = 'mymodel', model_config = 'ylcls/configs/ResNet/ResNet50.yaml'),
    # model_args = dict(model_type='timmmodel', model_config=dict(model_name='resnet50', pretrained=False, num_classes=102)),
    model = None,
    index2label_file_path = 'garbage/index2label.json',
    index2label = None,
    train_txt='train.txt',
    valid_txt='valid.txt',
    test_txt='test.txt',
    dataset_path = "garbage",
    dataset_name = 'CLSDataset',
    train_images_file_list = None,
    train_labels_list = None,
    valid_images_file_list = None,
    valid_labels_list = None,
    test_images_file_list = None,
    test_labels_list = None,
    train_dataset = None,
    valid_dataset = None,
    train_dataloader = None,
    valid_dataloader = None,
    cache = False,
    workers = 0,
    mosaic = False,
    transforms = None,
    max_nums = None
)

